#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# @Author: renjin@bit.edu.cn
# @Date  : 2025-11-02


import threading
from queue import Queue
import numpy as np
import time
import os
from typing import Union
from spirems import Subscriber, Publisher, cvimg2sms, sms2cvimg, def_msg, QoS, BaseNode, get_extra_args
from spirems.mod_helper import download_model
import argparse
import platform
import cv2
from sklearn.cluster import DBSCAN


def trans_det_results(x_min, y_min, x_max, y_max, w, h):
    sms_results = def_msg('spirecv_msgs::2DTargets')

    sms_results["file_name"] = ""
    sms_results["height"] = int(h)
    sms_results["width"] = int(w)
    sms_results["targets"] = []

    ann = dict()
    ann["category_name"] = "moving_object"
    ann["category_id"] = 0
    ann["score"] = 1.0
    ann["bbox"] = [float(x_min), float(y_min), float(x_max-x_min), float(y_max-y_min)]
    sms_results["targets"].append(ann)

    return sms_results


class AdaptiveThresholdNode(threading.Thread, BaseNode):
    def __init__(
        self,
        job_name: str,
        ip: str = '127.0.0.1',
        port: int = 9094,
        param_dict_or_file: Union[dict, str] = None,
        sms_shutdown: bool = True,
        **kwargs
    ):
        threading.Thread.__init__(self)
        sms_shutdown = True if sms_shutdown in ['True', 'true', '1', True] else False
        BaseNode.__init__(
            self,
            self.__class__.__name__,
            job_name,
            ip=ip,
            port=port,
            param_dict_or_file=param_dict_or_file,
            sms_shutdown=sms_shutdown,
            **kwargs
        )

        self.eps = self.get_param("eps", 10)
        self.min_samples = self.get_param("min_samples", 5)
        self.use_shm = self.get_param("use_shm", -1)
        self.params_help()

        self.b_use_shm = False
        if self.use_shm == 1 or (self.use_shm == -1 and platform.system() == 'Linux'):
            self.b_use_shm = True

        self._show_writer = Publisher(
            '/' + job_name + '/detector/image_results', 'memory_msgs::RawImage' if self.b_use_shm else 'sensor_msgs::CompressedImage',
            ip=ip, port=port
        )

        input_url = '/' + job_name + '/sensor/image_raw'
        output_url = '/' + job_name + '/detector/results'

        self.job_queue = Queue()
        self.queue_pool.append(self.job_queue)

        self._image_reader = Subscriber(
            input_url, 'std_msgs::Null', self.image_callback,
            ip=ip, port=port
        )
        self._result_writer = Publisher(
            output_url, 'spirecv_msgs::2DTargets',
            ip=ip, port=port, qos=QoS.Reliability
        )

        self.last_img = None
        self.last_img_name = None
        self.last_vid_name = None

        self.infer_n_ims = 0
        self.infer_time3 = 0.0
        self.infer_time2 = 0.0
        self.infer_time1 = 0.0

        self.start()

    def release(self):
        BaseNode.release(self)
        self._image_reader.kill()
        self._result_writer.kill()
        self._show_writer.kill()

    def image_callback(self, msg):
        if not self.job_queue.empty():
            self.job_queue.queue.clear()
        img = sms2cvimg(msg)
        self.job_queue.put({'msg': msg, 'img': img})

    def run(self):
        while self.is_running():
            msg_dict = self.job_queue.get(block=True)
            if msg_dict is None:
                break

            t1 = time.time()
            msg, img_bgr = msg_dict['msg'], msg_dict['img']
            img = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY)

            binary_img = cv2.adaptiveThreshold(img, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY, 15, 6)
            binary_img = 255 - binary_img

            # 提取所有255像素的坐标
            y_coords, x_coords = np.where(binary_img == 255)
            if len(x_coords) == 0:
                print("图像中没有255像素")
                exit()

            # 构建坐标点集（N个点，每个点为(x,y)）
            points = np.column_stack((x_coords, y_coords))

            # DBSCAN聚类（参数需根据实际数据调整）
            # eps：邻域半径，min_samples：形成簇的最小点数
            dbscan = DBSCAN(eps=self.eps, min_samples=self.min_samples)  # eps越小，聚类越密集
            labels = dbscan.fit_predict(points)

            # 排除噪声点（label=-1）
            unique_labels = np.unique(labels)
            valid_labels = [label for label in unique_labels if label != -1]

            if not valid_labels:
                print("未找到有效的聚类簇（所有点均为噪声）")
            else:
                # 找到点数最多的簇（最集中的区域）
                max_cluster_size = -1
                best_cluster = None
                for label in valid_labels:
                    cluster_points = points[labels == label]
                    if len(cluster_points) > max_cluster_size:
                        max_cluster_size = len(cluster_points)
                        best_cluster = cluster_points

                # 计算该簇的外包矩形
                x_min, y_min = np.min(best_cluster, axis=0).astype(int)
                x_max, y_max = np.max(best_cluster, axis=0).astype(int)

                # 绘制红色矩形
                # color_img = cv2.cvtColor(binary_img, cv2.COLOR_GRAY2BGR)
                # cv2.rectangle(color_img, (x_min, y_min), (x_max, y_max), (0, 0, 255), 2)
                res_msg = trans_det_results(x_min, y_min, x_max, y_max, binary_img.shape[1], binary_img.shape[0])
                self._result_writer.publish(res_msg)
                # cv2.imshow("img", color_img)
                # cv2.waitKey(5)
                if self.b_use_shm:
                    msg = self._show_writer.cvimg2sms_mem(img_bgr)
                msg['spirecv_msgs::2DTargets'] = res_msg
                self._show_writer.publish(msg)

            print(time.time() - t1)

        self.release()
        print('{} quit!'.format(self.__class__.__name__))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--config',
        type=str,
        default='default_params.json',
        help='SpireCV2 Config (.json)')
    parser.add_argument(
        '--job-name', '-j',
        type=str,
        default='live',
        help='SpireCV Job Name')
    parser.add_argument(
        '--ip',
        type=str,
        default='127.0.0.1',
        help='SpireMS Core IP')
    parser.add_argument(
        '--port',
        type=int,
        default=9094,
        help='SpireMS Core Port')
    # args = parser.parse_args()
    args, unknown_args = parser.parse_known_args()
    if not os.path.isabs(args.config):
        current_path = os.path.abspath(__file__)
        params_dir = os.path.join(current_path[:current_path.find('spirecv-pro') + 11], 'params', 'spirecv2')
        args.config = os.path.join(params_dir, args.config)
    print("--config:", args.config)
    print("--job-name:", args.job_name)
    extra = get_extra_args(unknown_args)

    node = AdaptiveThresholdNode(args.job_name, param_dict_or_file=args.config, ip=args.ip, port=args.port, **extra)
    node.spin()
